/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2021. All rights reserved.
 * Description: 集合通信算子图优化头文件
 * Author: lilianlin
 * Create: 2019-12-7
 */

#ifndef HVD_GRAPH_OPTIMIZER_H
#define HVD_GRAPH_OPTIMIZER_H

#include "common/optimizer/graph_optimizer.h"
#include "common/optimizer/graph_optimizer_types.h"
#include "graph/compute_graph.h"

#include "graph/debug/ge_attr_define.h"
#include "framework/common/types.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/node.h"
#include "hvd_ops_kernel_info_store.h"
#include "graph/debug/ge_attr_define.h"

namespace hccl {
const string HVD_GRAPH_OPTIMIZER_NAME = "hvd_graph_optimizer";
constexpr u32 SPLIT_STREAM_THREAD = 400; // each 400 op will assign a new stream
constexpr std::int64_t HCCL_HVD_FORMAT_PAIRED_INPUT_OUTPUT = 2; // 输入输出地格式须相同

class HvdGraphOptimizer : public ge::GraphOptimizer {
public:
    HvdGraphOptimizer();
    ~HvdGraphOptimizer() override;
    ge::Status Initialize(const map<std::string, std::string>& options,
                          ge::OptimizeUtility *const optimizeUtility) override;
    // close graphOptimizer
    ge::Status Finalize() override;
    // optimize original graph for FE quant optimize
    ge::Status OptimizeGraphPrepare(ge::ComputeGraph& graph) override;
    // optimize original graph, using in graph preparation stage
    ge::Status OptimizeOriginalGraph(ge::ComputeGraph& graph) override;
    // optimize fused graph
    ge::Status OptimizeFusedGraph(ge::ComputeGraph& graph) override;
    // optimize whole graph, using after graph merged stage
    ge::Status OptimizeWholeGraph(ge::ComputeGraph& graph) override;
    // get attribute of graph optimizer
    ge::Status GetAttributes(ge::GraphOptimizerAttribute& attrs) const override;

private:
    HcclResult CheckSupportedOP(const std::string &sCollectiveType) const;
    HcclResult AssignNewStream(u32 &opCount, ge::NodePtr &nodePtr, std::string &waitStream, std::string &sendStream);
    HcclResult AddControlAnchor(ge::NodePtr &nodePtr, ge::NodePtr &addedNodePtr);
    HcclResult AddDependence(std::vector<ge::NodePtr> &preDependNodes);
    ge::Status CalcOpRunningParam(ge::Node& node);
    HcclResult SetOpOutputMemSize(ge::Node& node, const std::string &sCollectiveType);
    HcclResult CalcHCCLOutputMemSize(const std::string &sCollectiveType, int64_t& memSize);
    HcclResult SetOpMemAttr(ge::Node& node, const std::string &sCollectiveType);
    HcclResult SetOpAtomicInputIndex(ge::Node& node, const std::string &sCollectiveType);
};
}

#endif
